﻿#ifndef __ROAD_SIGN_TRAINER_H__
#define __ROAD_SIGN_TRAINER_H__
#include "base64.h"

#include "ColorTrainer.h"
#include "SoftMaxTrainer.h"

class RoadSignTrainer : public SoftMaxTrainer
{
public:
    RoadSignTrainer()
    {
        m_max_epochs      = 4000;
        m_batch           = 8;
        m_auto_stop_train = false;

        CreateGetClassifyIndex(u8R"(左转)");
        CreateGetClassifyIndex(u8R"(右转)");
        CreateGetClassifyIndex(u8R"(直行)");
        CreateGetClassifyIndex(u8R"(停车)");
        CreateGetClassifyIndex(u8R"(隧道)");
        CreateGetClassifyIndex(u8R"(禁止通行)");
    }

    void LoadNet()
    {
        auto net     = GetNet();
        if (net->GetLayers()->size() > 2) 
        {
            RefreshSoftMaxLayer();
            return;
        }
        auto options = std::make_shared<std::vector<std::shared_ptr<Options>>>();
        auto option  = std::make_shared<Options>("net");
        option->SetOption("width", 100);
        option->SetOption("height", 60);
        option->SetOption("channels", 3);
        options->push_back(option);

        option = std::make_shared<Options>("convolutional");
        option->SetOption("size", 3);
        option->SetOption("stride", 2);
        option->SetOption("filters", 16);
        options->push_back(option);

        option = std::make_shared<Options>("convolutional");
        option->SetOption("size", 3);
        option->SetOption("stride", 2);
        option->SetOption("filters", 32);
        options->push_back(option);

        option = std::make_shared<Options>("convolutional");
        option->SetOption("size", 3);
        option->SetOption("stride", 2);
        option->SetOption("filters", 16);
        options->push_back(option);

        option = std::make_shared<Options>("connected");
        option->SetOption("output", GetClassifySize());
        option->SetOption("activation", "linear");
        options->push_back(option);

        option = std::make_shared<Options>("softmax");
        options->push_back(option);

        net->Load(options);
    }

    bool LoadModelFromFile(const std::string &model_path)
    {
        return Trainer::LoadModelFromFile(model_path);
    }

    std::string PredictBase64Image(std::string img_data)
    {
        return PredictBinaryImage(base64_decode(img_data));
    }

    std::string PredictBinaryImage(std::string img_data)
    {
        auto img = std::make_shared<DataImage>();
        if (!img->Load((unsigned char *)(img_data.c_str()), img_data.size()))
        {
            std::cout << "invalid image data!!!" << std::endl;
            return "";
        }
        return Predict(img);
    }

    std::string PredictPathImage(const std::string &img_path)
    {
        auto img = std::make_shared<DataImage>();
        if (!img->Load(img_path))
        {
            std::cout << "invalid image path: " << img_path << std::endl;
            return "";
        }
        return PredictImage(img);
    }

    std::string PredictImage(std::shared_ptr<DataImage> img)
    {
        return Predict(img == nullptr ? nullptr : GetColorTrainer()->GetColorImage(img, GetRoadSignColorClassifyIndex(), GetRoadSignSize()));
    }

    void Train(const std::string& train_path)
    {
        auto net = GetNet();
        if (net->GetLayers()->size() == 0)
        {
            LoadNet();
        }
        ResetDataSet(GetDataSet());
        if (!train_path.empty()) 
        {
            GetDataSet()->Load(train_path);
        }
        Trainer::Train();
    }

    void LoadDataItems()
    {
        ClearClassifyIndex();

        for (auto &item : m_path_classify_map)
        {
            CreateGetClassifyIndex(item.second);
        }

        auto dataitems     = GetDataSet()->GetDataItems();
        auto classify_size = GetClassifySize();

        for (auto &item : m_path_classify_map)
        {
            std::string image_path = item.first;
            int classify_index     = CreateGetClassifyIndex(item.second);

            auto img = std::make_shared<DataImage>();
            if (std::filesystem::path(image_path).extension() == ".base64")
            {
                std::string jpeg_prefix     = "data:image/jpeg;base64,";
                std::string png_prefix      = "data:image/png;base64,";
                std::string base64_img_data = LoadFile(image_path);
                if (base64_img_data.substr(0, jpeg_prefix.length()) == jpeg_prefix)
                {
                    base64_img_data = base64_img_data.substr(jpeg_prefix.length());
                }
                if (base64_img_data.substr(0, png_prefix.length()) == png_prefix)
                {
                    base64_img_data = base64_img_data.substr(png_prefix.length());
                }
                auto img_data = base64_decode(base64_img_data);
                if (!img->Load((const unsigned char *)(img_data.data()), img_data.size()))
                {
                    std::cout << "invalid image data: " << image_path << std::endl;
                    continue;
                }
            }
            else
            {
                if (!img->Load(image_path))
                {
                    std::cout << "invalid image path: " << image_path << std::endl;
                    continue;
                }
            }

            img = GetColorTrainer()->GetColorImage(img, GetRoadSignColorClassifyIndex(), GetRoadSignSize());
            if (img == nullptr)
            {
                std::cout << "不存在红色路标!!!" << image_path << std::endl;
                continue;
            }

            auto dataitem                         = std::make_shared<DataItem>();
            auto truth                            = std::make_shared<Truth>(classify_size);
            *(truth->GetDatas() + classify_index) = 1.f;

            dataitem->SetImage(img);
            dataitem->GetTruths()->push_back(truth);
            dataitem->SetLoaded(true);
            dataitems->push_back(dataitem);
        }
    }

public:
    inline ColorTrainer *GetColorTrainer() { return ColorTrainer::GetInstance(); }
    inline int GetRoadSignColorClassifyIndex() { return GetColorTrainer()->CreateGetClassifyIndex("red"); }
    inline int GetRoadSignSize() { return 160; }
};

#endif